package edu.zju.gis.dbfg.model.loader;

import com.google.gson.Gson;
import org.apache.http.HttpHost;
import org.apache.http.client.config.RequestConfig;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.PairFlatMapFunction;
import org.elasticsearch.action.admin.indices.settings.put.UpdateSettingsRequest;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.CreateIndexResponse;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.IOException;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.*;

/**
 * @author Hu
 * @date 2019/10/6
 * ES 数据入库基类
 **/
public abstract class ESDataLoader2<T> extends DataLoader<T> {

    private final static Logger logger = LoggerFactory.getLogger(ESDataLoader2.class);

    protected static int RETRY_TIME = 3;
    protected static int NUMBER_OF_SHARDS = 3;
    protected static int NUMBER_OF_REPLICA = 1;
    protected static int INDEX_REFRESH_INTERVAL = -1;
    protected static int INDEX_NORMAL_REFRESH_INTERVAL = 2;
    protected static int INDEX_TIME_OUT = 5; // time unit is minute
    protected static String ES_IP = "10.2.40.51,10.2.40.52,10.2.40.53,10.2.40.54";
    protected static int ES_PORT = 9200;
    protected final static Integer STEP_SIZE = 1000;

    private String indexName = null;

    protected abstract JavaRDD<T> readToRDD(JavaSparkContext jsc);

    protected BulkResponse batchLoad(RestHighLevelClient client, String name, List<T> ts, int count) throws IOException {
        if (count >= RETRY_TIME) {
            logger.error("Retry time: " + count + ", Abort");
            return null;
        } else if (count > 1) {
            logger.info("Retry time:" + count);
        }
        BulkRequest br = new BulkRequest();
        br.timeout(TimeValue.timeValueMinutes(INDEX_TIME_OUT));
        try {
            Method m = getTClass().getDeclaredMethod("getId");
            Method mg = getTClass().getDeclaredMethod("toGeoJson");
            for(T t:ts) {
                IndexRequest ir = new IndexRequest(name);
                String id = (String) m.invoke(t);
                if (id == null || id.trim().length() == 0) {
                    id = UUID.randomUUID().toString();
                }
                ir.id(id);
                String geoJson = (String) mg.invoke(t);
                ir.source(geoJson, XContentType.JSON);
                br.add(ir);
            }
            return client.bulk(br, RequestOptions.DEFAULT);
        } catch (Exception e) {
            e.printStackTrace();
            String msg = e.getMessage();
            logger.error(msg);
            int z = count + 1;
            return batchLoad(client, name, ts, z);
        }
    }

    protected void createIndex(String name) throws IOException {

        String ip = this.ES_IP.split(",")[0].trim();

        this.indexName = name;

        RestHighLevelClient client = new RestHighLevelClient(
                RestClient.builder(new HttpHost(ip, this.ES_PORT)));

        logger.info("Create index: " + name);

        // index exist
        GetIndexRequest indexExistRequst = new GetIndexRequest(name);
        boolean exists = client.indices().exists(indexExistRequst, RequestOptions.DEFAULT);
        if (exists) {
            logger.warn(String.format("Index already exists: %s", name));
            return;
        }

        CreateIndexRequest createIndexRequest = new CreateIndexRequest(name);
        createIndexRequest.settings(Settings.builder()
                .put("index.number_of_shards", this.NUMBER_OF_SHARDS)
                .put("index.number_of_replicas", this.NUMBER_OF_REPLICA)
                .put("index.refresh_interval", this.INDEX_REFRESH_INTERVAL)
        );

        // get es mapping of class T
        List<Field> fields = this.getFields();
        Map<String, Object> properties = new HashMap<>();
        for(Field field: fields) {
            Map<String, String> fieldMap = new HashMap<>();
            if (field.getType().getName().equals(Integer.class.getName())) {
                fieldMap.put("type", "integer");
            } else if (field.getType().getName().equals(Double.class.getName())) {
                fieldMap.put("type", "double");
            } else if (field.getType().getName().equals(String.class.getName())) {
                fieldMap.put("type", "text");
                // TODO 默认 ik 分词，之后改成通过 field annotation 指定
//        fieldMap.put("analyzer", "ik_max_word");
//        fieldMap.put("search_analyzer", "ik_smart");
            } else if (field.getType().getName().equals(Geometry.class.getName())){
                fieldMap.put("type", "geo_shape");
            } else {
                fieldMap.put("type", "text");
            }
            properties.put(field.getName(), fieldMap);
        }
        Map<String, Object> mapping = new HashMap<>();
        mapping.put("properties", properties);
        Gson gson = new Gson();
        createIndexRequest.mapping(gson.toJson(mapping), XContentType.JSON);

        CreateIndexResponse createIndexResponse = client.indices().create(createIndexRequest, RequestOptions.DEFAULT);
        if (createIndexResponse.isAcknowledged()) {
            logger.info(String.format("create index %s success", name));
        } else {
            logger.error(String.format("create index %s fail", name));
            System.exit(-1);
        }
        client.close();
    }


    protected void indexDocs(String name) throws IOException {

        String ip = this.ES_IP.split(",")[0].trim();

        RestHighLevelClient client = new RestHighLevelClient(
                RestClient.builder(new HttpHost(ip, this.ES_PORT)));

        logger.info("Update index configuration and index documents: " + name);

        // index exist
        GetIndexRequest indexExistRequst = new GetIndexRequest(name);
        boolean exists = client.indices().exists(indexExistRequst, RequestOptions.DEFAULT);
        if (!exists) {
            logger.warn(String.format("Index not exists: %s", name));
            return;
        }

        UpdateSettingsRequest updateIndexRequest = new UpdateSettingsRequest(name);
        updateIndexRequest.settings(Settings.builder()
                .put("index.refresh_interval", this.INDEX_NORMAL_REFRESH_INTERVAL));

        AcknowledgedResponse response = client.indices().putSettings(updateIndexRequest, RequestOptions.DEFAULT);
        if (response.isAcknowledged()) {
            logger.info(String.format("update index %s success", name));
        } else {
            logger.error(String.format("update index %s fail", name));
            System.exit(-1);
        }

        client.close();
    }

    @Override
    protected IngestResult ingest (JavaRDD<T> featuresRDD, String sourceName) throws Exception {
        JavaPairRDD<String, Object> ingestResultRDD = featuresRDD.mapPartitionsToPair(new PairFlatMapFunction<Iterator<T>, String, Object>() {
            @Override
            public Iterator<Tuple2<String, Object>> call(Iterator<T> reader) throws Exception {
                long total = 0;
                List<T> ts = new ArrayList<>();
                List<Tuple2<String, Object>> result = new ArrayList<>();
                String[] ips = ES_IP.split(",");
                HttpHost[] hosts = new HttpHost[ips.length];
                for (int m=0; m<ips.length; m++) {
                    hosts[m] = new HttpHost(ips[m].trim(), ES_PORT);
                }
                RestHighLevelClient clientT = new RestHighLevelClient(
                        RestClient.builder(hosts).setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
                            @Override
                            public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder builder) {
                                return builder.setConnectTimeout(INDEX_TIME_OUT * 60 * 1000).setSocketTimeout(INDEX_TIME_OUT * 60 *1000);
                            }
                        }));
                while(reader.hasNext()) {
                    ts.add(reader.next());
                    if (ts.size() >= STEP_SIZE) {
                        total += ts.size();
                        BulkResponse br1 = batchLoad(clientT, sourceName, ts, 1);
                        if (br1 == null) {
                            logger.warn("Abort: " + ts.size());
                            result.addAll(processAborteResult(ts));
                        } else {
                            result.addAll(processBatchLoadResponse(ts, br1));
                        }
                        ts.clear();
                    }
                }

                // 对剩余的 feature 进行入库操作
                if (ts.size() > 0) {
                    total += ts.size();
                    BulkResponse br2 = batchLoad(clientT, sourceName, ts, 1);
                    if (br2 == null) {
                        logger.warn("Abort: " + ts.size());
                        result.addAll(processAborteResult(ts));
                    } else {
                        result.addAll(processBatchLoadResponse(ts, br2));
                    }
                }

                result.add(new Tuple2<>("total", total));
                clientT.close();
                return result.iterator();
            }
        });
        return this.processIngestResult(ingestResultRDD);
    }

    @Override
    protected void calExtents(JavaRDD<T> featuresRDD) {
        JavaPairRDD<String, Double> extendsRDD = featuresRDD.flatMapToPair(new PairFlatMapFunction<T, String, Double>() {
            @Override
            public Iterator<Tuple2<String, Double>> call(T t) throws Exception {
                Method method = getTClass().getDeclaredMethod("getShape");
                Geometry g = (Geometry) method.invoke(t);
                List<Tuple2<String, Double>> result = new ArrayList<>();
                Envelope envelope = g.getEnvelopeInternal();
                result.add(new Tuple2<>("xMin", envelope.getMinX()));
                result.add(new Tuple2<>("xMax", -1 * envelope.getMaxX()));
                result.add(new Tuple2<>("yMin", envelope.getMinY()));
                result.add(new Tuple2<>("yMax", -1 * envelope.getMaxY()));
                return result.iterator();
            }
        }).reduceByKey((a1, a2)->a1 <= a2 ? a1: a2);
        this.processCalExtentResult(extendsRDD);
    }

    protected List<Tuple2<String, Object>> processAborteResult(List<T> ts) {
        List<Tuple2<String, Object>> result = new ArrayList<>();
        for (T t: ts) {
            result.add(new Tuple2<String, Object>("abort", t));
        }
        return result;
    }

    protected List<Tuple2<String, Object>> processBatchLoadResponse(List<T> ts, BulkResponse bulkResponse){
        List<Tuple2<String, Object>> result = new ArrayList<>();
        if (bulkResponse.hasFailures()) {
            int i=0;
            for (BulkItemResponse bulkItemResponse : bulkResponse) {
                if (bulkItemResponse.isFailed()) {
                    BulkItemResponse.Failure failure =
                            bulkItemResponse.getFailure();
                    T t = ts.get(i);
                    result.add(new Tuple2<>("error_msg", String.format(":%s", failure.getMessage())));
                    result.add(new Tuple2<>("error_obj", t));
                }
                i++;
            }
        }
        return result;
    }

    protected IngestResult processIngestResult(JavaPairRDD<String, Object> ingestResultRDD) throws Exception {
        List<String> errors = new ArrayList<>();
        List<String> aborts = new ArrayList<>();
        List<T> errorObjs = new ArrayList<>();
        long to = 0;
        List<Tuple2<String, Object>> ingestResult = ingestResultRDD.collect();
        Method mg = getTClass().getDeclaredMethod("toGeoJson");
        for(Tuple2<String, Object> r: ingestResult) {
            if (r._1.equals("error_msg")){
                errors.add((String) r._2);
            } else if (r._1.equals("total")) {
                to += (Long) r._2;
            } else if (r._1.equals("abort")) {
                String geoJson = (String) mg.invoke(r._2);
                aborts.add(geoJson);
            } else if(r._1.equals("error_obj")) {
                errorObjs.add((T) r._2);
            }
        }

        logger.info("Loader finish success");
        logger.info(String.format("total = %d", to));
        logger.info(String.format("error_count = %d", errors.size()));
        logger.info("======== Errors =======");
        errors.forEach(logger::info);
        logger.info("======== Error Objects =======");
        errorObjs.forEach(x->logger.info(x.toString()));
        logger.info("======== Abort =======");
        aborts.forEach(logger::info);

        return new IngestResult((int)to, errors, aborts, errorObjs);
    }

    protected void processCalExtentResult(JavaPairRDD<String, Double> extendsRDD) {
        Map<String, Double> extendsMap = extendsRDD.collectAsMap();
        logger.info(String.format("xMin = %.6f", extendsMap.get("xMin")));
        logger.info(String.format("yMin = %.6f", extendsMap.get("yMin")));
        logger.info(String.format("xMax = %.6f", -1 * extendsMap.get("xMax")));
        logger.info(String.format("yMax = %.6f", -1 * extendsMap.get("yMax")));
    }

    @Override
    protected void finish() throws IOException {
        if (indexName != null) {
            indexDocs(indexName);
        }
        logger.info("Loader finish");
    }

    public class IngestResult {

        private int total;
        private List<String> errors;
        private List<String> aborts;
        private List<T> errorObj;

        public IngestResult(int total, List<String> errors, List<String> aborts, List<T> errorObj) {
            this.total = total;
            this.errors = errors;
            this.aborts = aborts;
            this.errorObj = errorObj;
        }

        public int getTotal() {
            return total;
        }

        public List<String> getErrors() {
            return errors;
        }

        public List<String> getAborts() {
            return aborts;
        }

        public List<T> getErrorObj() {
            return errorObj;
        }

    }

}
